# -*- coding: utf-8 -*-
# @Time    : 2021/9/29 17:25
# @Author  : huangwei
# @File    : get_table_data.py
# @Software: PyCharm
import cv2
import os
from config import args
from method import create_dir
from signature import sort_box_to_group, cut_name
from table import get_cell
from table_net import table_net
from word import TextDetector, TextRecognizer


def get_table_data( filepath, text_detector, tableNet ):
    img = cv2.imread(filepath)

    # 创建临时文件夹
    tmp_file_name = os.path.basename(filepath)
    file_name, _ = os.path.splitext(tmp_file_name)
    temp_dir = "temp_path/{}".format(file_name)
    create_dir(temp_dir)

    # 二值化处理
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    _, thresh = cv2.threshold(gray, 160, 255, cv2.THRESH_BINARY)
    new_img = cv2.cvtColor(thresh, cv2.COLOR_GRAY2BGR)
    new_img_path = "{0}/{1}".format(temp_dir, tmp_file_name)
    cv2.imwrite(new_img_path, new_img)

    # 检测文字框
    dt_boxes = text_detector(img)
    print("文本框的数量为：", len(dt_boxes))

    # 获取表格线和表格框

    prob = 0.5  # 是否连通的概率判断
    row = 20  # 短于这个长度的横线舍去
    alph = 20  # 两条线之间的最大间隔，小于则连接两条线
    col = 20
    boxes, row_lines = get_cell(new_img_path, tableNet, prob, alph, row, col)

    # 将检测出的框画在图片上
    from method import draw_box
    for box in boxes:
        draw_box(img, box)
    cv2.imwrite("{}/box.png".format(temp_dir), img)
    print("检测出的表格框数量为：", len(boxes))

    # 使用检测出的表格线划分每一行检测出的表格和文字格
    group_boxes, group_det_boxes = sort_box_to_group(boxes, dt_boxes, row_lines)

    img3 = cv2.imread(filepath)
    data_list = []
    # # 对每一行的box进行识别和剪切
    for i in range(len(group_boxes)):
        cut_name(img3, group_boxes[i], group_det_boxes[i], i, text_recognizer, data_list)

    output_path = "{0}/out_{1}".format(temp_dir, tmp_file_name)
    cv2.imwrite(output_path, img3)
    print(data_list)
    # print(len(data_list))


img_path = "images/simple3.png"

### 加载模型
text_detector = TextDetector(args)
text_recognizer = TextRecognizer(args)
# WEIGHT_PATH = "models/table.weights"
# tableNet = cv2.dnn.readNetFromDarknet(WEIGHT_PATH.replace('.weights', '.cfg'), WEIGHT_PATH)
tableModeLinePath = 'models/table-line.h5'
model = table_net((None, None, 3), 2)
model.load_weights(tableModeLinePath)

get_table_data(img_path, text_detector, model)
